import cv2
import numpy as np
# from matplotlib import pyplot as plt
import time
import math
import glob
import argparse


parser = argparse.ArgumentParser(prog='stitch20190602.py', description='Rotation model images stitcher')
parser.add_argument('img_names', nargs='+',help='files to stitch',type=str)
parser.add_argument('--output',action = 'store', default = 'result.jpg',help='The default is "result.jpg"',type=str,dest = 'output' )
#截除多余部分
#warpImg 输入图像
#max     空白最大行数
def cutblank(warpImg ,max):
    mark=0
    break_flage =False
    warprows, warpcols = warpImg.shape[:2]
    if(warprows<0):
        print("warprows=",warprows)
    if (warpcols < 0):
        print("warpcols=", warpcols)
    for row in range((int)(warprows/2), warprows):
        for col in range(0, warpcols):
            if not warpImg[row, col].any():
                mark += 1  # 有标记加1
            if mark > 1:
                if warpImg[row, col].any():
                    mark -= 1  # 无，标记减1
            if mark > max:
                warpcols = col-max  # 连续20个像素无，认为后面无，截断
                break_flage=True
                break
        if break_flage==True:
            break

    res = np.zeros([warprows, warpcols, 3], np.uint8)
    for row in range(0, warprows):
        for col in range(0, warpcols):
            res[row, col][0] = warpImg[row, col][0]
            res[row, col][1] = warpImg[row, col][1]
            res[row, col][2] = warpImg[row, col][2]

    if mark<max:
        return warpImg
    else:
        return res


#柱面投影
#输入参数
#img1 输入图像时
#f图像焦距
def cylindricalWarpImage(img1, f, savefig=False):
    # w=img1.shape[1]
    # f = (w / 2) / math.atan(math.pi/ 8)
    # K = np.array([[f, 0, w/2], [0, f, h/2], [0, 0, 1]]) # mock calibration matrix  模板校准矩阵
    res = np.zeros_like(img1)
    cyl_h,cyl_w= res.shape[:2]
    x_c = float(cyl_w) / 2.0
    y_c = float(cyl_h) / 2.0
    leng = (int)(f * math.atan((0 - cyl_w / 2) / f) + cyl_w / 2)
    cyl = np.zeros([cyl_h, (cyl_w-2*leng),3], np.uint8)

    for x_cyl in np.arange(0,cyl_w):
        for y_cyl in np.arange(0,cyl_h):
            x1= f*math.atan((x_cyl-x_c)/f)+f*math.atan(x_c/f)
            y1= f*(y_cyl-y_c)/math.sqrt((x_cyl-x_c)*(x_cyl-x_c)+f*f)+y_c

            x_im=(int)(x1+0.5) #加0.5是为了四舍五入
            y_im=(int)(y1+0.5)
            if(x_im<cyl_w and y_im<cyl_h):
                res[int(y_im),int(x_im)][0]= img1[int(y_cyl),int(x_cyl)][0]
                res[int(y_im), int(x_im)][1] = img1[int(y_cyl), int(x_cyl)][1]
                res[int(y_im), int(x_im)][2] = img1[int(y_cyl), int(x_cyl)][2]
    if savefig:
        plt.imshow(cyl, cmap='gray')
        plt.savefig("cyl.png",bbox_inches='tight')

    return res


#img1 拼接的第一张图
#img2 拼接的第二张图

def stitch2img(img1,img2):
    MIN = 10
    MAXHIGH=800
    w1=img1.shape[1]
    w2 = img2.shape[1]
    h1=img1.shape[0]
    h2 = img2.shape[0]
    f1 = (w1 / 2) / math.atan(math.pi/ 9)
    f2= (w1 / 2) / math.atan(math.pi / 9)
    f=(int)((f1+f2)/2)

    if(h1>MAXHIGH): #如果图片大于设定值，按比例缩小图片
        resize_scale1=MAXHIGH/h1
        img1 = cv2.resize(src=img1, dsize=None, fx=resize_scale1, fy=resize_scale1, interpolation=cv2. INTER_LINEAR)
    if (h2 > MAXHIGH):  # 如果图片大于设定值，按比例缩小图片
        resize_scale2 = MAXHIGH / h2
        img2 = cv2.resize(src=img2, dsize=None, fx=resize_scale2, fy=resize_scale2, interpolation=cv2.INTER_LINEAR)



    # len2 = (int)(f2 * math.atan((0 - w2/2) / f2) + w2/2)

    img1 =cylindricalWarpImage(img1,f)
    img2 =cylindricalWarpImage(img2,f)
    surf = cv2.xfeatures2d.SURF_create(400)
    kp1, descrip1 = surf.detectAndCompute(img1, None)
    kp2, descrip2 = surf.detectAndCompute(img2, None)


    FLANN_INDEX_KDTREE = 0
    indexParams = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    searchParams = dict(checks=50)
    #单应性匹配
    flann = cv2.FlannBasedMatcher(indexParams, searchParams)
    match = flann.knnMatch(descrip1, descrip2, k=2)

    good = []
    for i, (m, n) in enumerate(match):
        if (m.distance < 0.75 * n.distance):
            good.append(m)

    if len(good) > MIN:
        # 相机内参
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        ano_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
        # 单应矩阵
        M, mask = cv2.findHomography(src_pts, ano_pts, cv2.RANSAC, 7.0)
        #cv2.warpPepective  就单应矩阵（单应矩阵反映两图的相互关系）
        #np.linalg.inv(M) 矩阵求逆
        #shape[0] 高度
        #shape[1] 宽度
        # The function warpPerspective transforms the source image using the specified matrix:
        #根据单应矩阵 以图2为基础创建拼接掩膜
    warpImg = cv2.warpPerspective(img2, np.linalg.inv(M), (img1.shape[1] + (int)(img2.shape[1]), img1.shape[0]))


    direct = warpImg.copy()
    #0-图1高（行），0-图1寛（列）为图1
    #在掩膜左边放入图1
    direct[0:img1.shape[0], 0:img1.shape[1]] = img1


        # cv2.imshow("warpImg",warpImg)
        # cv2.waitKey(0)
        # 取图1的行和列

    #融合处理
    rows, cols = img1.shape[:2]

    for col in range(0, cols):
        if img1[:, col].any() and warpImg[:, col].any():  # 开始重叠的最左端
            left = col
            break

    for col in range(cols - 1, 0, -1):
        if img1[:, col].any() and warpImg[:, col].any():  # 重叠的最右一列
            right = col
            break
    #定义融合图的矩阵
    res = np.zeros([rows, cols, 3], np.uint8)
    #遍历整个融合图像
    for row in range(0, rows):
        for col in range(0, cols):
            if not img1[row, col].any():  # 如果图1没有原图，用图2的
                res[row, col] = warpImg[row, col]
            elif not warpImg[row, col].any():#如果图2没有原图，用图1的
                res[row, col] = img1[row, col]
            else:  #如果图1图2都没有原图，用图1图2融合处按权重算出来的像素
                srcImgLen = float(abs(col - left))
                testImgLen = float(abs(col - right))
                alpha = srcImgLen / (srcImgLen + testImgLen)
                #按权重叠加重叠部分，向着左图方向越远离衔接处的左图权重越大，右图的权重越小
                res[row, col] = np.clip(img1[row, col] *(1 - alpha)+ warpImg[row, col] * alpha  , 0, 255)

    warpImg[0:img1.shape[0], 0:img1.shape[1]] = res
    # if(img1.shape[0]>0 and img1.shape[1]>0 ):
    #     warpImg = cutblank(warpImg, 5) #切除多余的黑框
    return (warpImg)



def main():
    images = []
    num_images = 0
    args = parser.parse_args()
    img_names=args.img_names
    print(img_names)
    # stitchmoreimg2(img_names)

    startime = time.time()
    for name in img_names:
        # 读取图片
        img = cv2.imread(name)
        num_images += 1  # 照片计数
        images.append(img)
    if num_images < 2:
        print("need more image!")

    for i in range(num_images):
        if i == 0:
            reslt = stitch2img(images[i], images[i + 1])
        else:
            reslt = stitch2img(reslt, images[i])

    reslt= cutblank(reslt, 20)
    endtime = time.time()
    print(endtime-startime)
    cv2.imshow("reslt",reslt)
    cv2.imwrite("reslt.jpg", reslt)
    cv2.waitKey()


if __name__ == '__main__':
    main()
    cv2.destroyAllWindows()